import os

import pandas as pd
from pypinyin import lazy_pinyin
from tools import *
from tools.translator import Translator
# from random import randint
# from random import random
import random
import re
from bdtime import Time


def get_one_random_row_index_from_df(df):
    r_i = random.randint(0, df.shape[0] - 1)
    return r_i


class DataMaker:
    def __init__(self,
                 p_of_add_prefix_of_title=0.5,
                 name_sample_range=(2, 5),
                 institution_sample_range=(1, 5),
                 ):
        self.p_of_add_prefix_of_title = p_of_add_prefix_of_title
        self.name_sample_range = name_sample_range
        self.institution_sample_range = institution_sample_range

        self.columns = ['text', 'label', 'label_span']
        self.df = pd.DataFrame([], columns=self.columns)

        self.name_df: pd.DataFrame = pd.read_excel('my_data/names.xlsx')
        self.institution_df: pd.DataFrame = pd.read_excel('my_data/institutions.xlsx')
        self.article_df: pd.DataFrame = pd.read_excel('my_data/articles.xlsx')

        self.labels = self.get_labels()

    def get_labels(self):
        _labels = ['title', 'authors', 'institutions']
        _labels_1 = ['cn_' + i for i in _labels] + ['en_' + i for i in _labels]

        labels = ['O']
        for i in range(len(_labels_1)):
            _labels_i = _labels_1[i]
            labels.append(f'B_{_labels_i}')
            labels.append(f'I_{_labels_i}')
        return labels

    _translator = None

    @property
    def translator(self):
        if not self._translator:
            print('--- init Translator!')
            self._translator = Translator()
        return self._translator

    def _get_prefix_of_title(self) -> str:
        prefix_of_title = ""
        if random.random() < self.p_of_add_prefix_of_title:
            r_i = get_one_random_row_index_from_df(self.article_df)
            content = self.article_df.loc[r_i, '文本全文']

            rd = random.randint(30, 100)
            prefix_of_title = content[-rd:]
        return prefix_of_title

    def _get_title_and_content(self) -> (str, str):
        r_i = get_one_random_row_index_from_df(self.article_df)
        row = self.article_df.loc[r_i, :]
        cn_title = row['标题']
        content = row['文本全文']
        en_title = self.translator.translate_cn_to_en(cn_title)[0]
        return cn_title, en_title, content

    def _get_author_ls(self) -> (list, list):
        sample_n = random.randint(*self.name_sample_range)
        sampled_df = self.name_df.sample(n=sample_n)
        name_zh_ls = sampled_df['name'].tolist()
        name_en_ls = [conv_name_from_zh_to_en(name_zh) for name_zh in name_zh_ls]
        return name_zh_ls, name_en_ls

    def _get_authors(self, name_zh_ls, name_en_ls, length_of_institutions) -> (str, str):
        if length_of_institutions == 1:
            cn_authors = '，'.join(name_zh_ls)
            en_authors = ','.join(name_en_ls)
        else:
            en_authors = cn_authors = ""
            for i in range(len(name_zh_ls)):
                if random.random() < 0.5:
                    len_i = 1
                else:
                    len_i = random.randint(1, length_of_institutions)
                institution_id_ls = [random.randint(1, length_of_institutions) for j in range(len_i)]
                institution_id_ls = sorted(set(institution_id_ls))
                institution_id_str = ','.join([str(_) for _ in institution_id_ls])
                cn_i = name_zh_ls[i]
                en_i = name_en_ls[i]
                suffix = ', ' if i != len(name_zh_ls) - 1 else ''
                cn_authors += cn_i + institution_id_str + suffix
                en_authors += en_i + institution_id_str + suffix
        return cn_authors, en_authors

    def _get_institution_ls(self) -> (list, list):
        sample_n = random.randint(*self.institution_sample_range)
        sampled_df = self.institution_df.sample(n=sample_n)
        cn_institution_ls = sampled_df['中文机构名'].tolist()
        en_institution_ls = sampled_df['英文机构'].tolist()
        cn_institution_ls = [i.strip(";").strip('；').strip('.') for i in cn_institution_ls]
        en_institution_ls = [i.strip(";").strip('；').strip('.') for i in en_institution_ls]
        return cn_institution_ls, en_institution_ls

    def _get_institutions(self, cn_institution_ls, en_institution_ls) -> (str, str):
        dtype = {
            '本节点编码': str,
            '父节点编码': str,
        }
        self.area_df: pd.DataFrame = pd.read_excel('my_data/地区树节点表.xlsx', dtype=dtype)
        self.city_df = self.area_df[self.area_df['节点等级'] == 2]

        length_of_institutions = len(cn_institution_ls)

        sampled_df = self.city_df.sample(length_of_institutions)

        def get_parent_name(row):
            # row = sampled_df.iloc[0]
            parent_code = row['父节点编码']
            parent_row = self.area_df[self.area_df['本节点编码'] == parent_code]
            assert not parent_row.empty and parent_row.shape[0] == 1
            parent_name = parent_row['节点名'].tolist()[0]
            return parent_name
        sampled_df['父节点名'] = sampled_df.apply(get_parent_name, axis=1)

        city_name_ls = sampled_df['节点名'].tolist()
        province_name_ls = sampled_df['父节点名'].tolist()
        # city_code_ls = sampled_df['本节点编码'].tolist()
        en_city_name_ls = self.translator.translate_cn_to_en(city_name_ls)
        en_province_name_ls = self.translator.translate_cn_to_en(province_name_ls)

        sampled_df['en_city_name_ls'] = en_city_name_ls
        sampled_df['en_province_name_ls'] = en_province_name_ls
        sampled_df['cn_institution_ls'] = cn_institution_ls
        sampled_df['en_institution_ls'] = en_institution_ls

        cn_str_ls, en_str_ls = [], []

        def get_str(row):
            cn_institution_i = row['cn_institution_ls']
            en_institution_i = row['en_institution_ls']

            city_code = row['本节点编码']
            city_name = row['父节点名']
            province_name = row['父节点名']

            en_city_name = row['en_city_name_ls'].strip('.')
            en_province_name = row['en_province_name_ls'].strip('.')

            cn_str_i = f'{cn_institution_i}, {province_name} {city_name}, {city_code}'
            en_str_i = f'{en_institution_i}, {en_province_name} {en_city_name}, {city_code}'

            cn_str_ls.append(cn_str_i)
            en_str_ls.append(en_str_i)

        sampled_df.apply(get_str, axis=1)

        cn_institutions = ""
        en_institutions = ""
        for i in range(length_of_institutions):
            suffix = '\n' if i != length_of_institutions - 1 else ''
            cn_institutions += f'{i+1}.' + cn_str_ls[i] + suffix
            en_institutions += f'{i+1}.' + en_str_ls[i] + suffix

        return cn_institutions, en_institutions

    def _get_summary(self, content):
        length = random.randint(30, 100)
        suffix = "..." if random.random() < 0.5 else ""
        cn_summary = content[:length] + suffix
        en_summary = self.translator.translate_cn_to_en(cn_summary)[0]
        return cn_summary, en_summary

    def _get_content(self):
        content = 1
        return content

    def get_context_dc(self):
        # debug = True
        debug = False

        tt = Time()
        if debug: print('======================\n')

        prefix_of_title = self._get_prefix_of_title()
        if debug: print('--- _get_prefix_of_title ------', tt.now())

        cn_institution_ls, en_institution_ls = self._get_institution_ls()
        cn_institutions, en_institutions = self._get_institutions(cn_institution_ls, en_institution_ls)
        length_of_institutions = len(cn_institution_ls)
        if debug: print('--- _get_institution_ls ------', tt.now())

        name_zh_ls, name_en_ls = self._get_author_ls()
        cn_authors, en_authors = self._get_authors(name_zh_ls, name_en_ls, length_of_institutions)
        if debug: print('--- _get_authors ------', tt.now())

        cn_title, en_title, content = self._get_title_and_content()
        if debug: print('--- _get_title_and_content ------', tt.now())

        # content = self._get_content()
        cn_summary, en_summary = self._get_summary(content)
        context_dc = {
            'prefix_of_title': prefix_of_title,
            'cn_info': {
                'title': cn_title,
                'authors': cn_authors,
                'institutions': cn_institutions,
                'summary': cn_summary,
            },
            'en_info': {
                'title': en_title,
                'authors': en_authors,
                'institutions': en_institutions,
                'summary': en_summary,
            },
            'content': content,
        }
        return context_dc

    def get_label_i_and_span(self, replaced_template_text, context_dc):
        labels = self.labels
        target_ls = labels[1:]

        label_i = [labels.index('O')] * len(replaced_template_text)
        label_span = {}
        for i in range(0, len(target_ls), 2):
            target_i = target_ls[i]

            target_label_0 = target_ls[i]
            target_label_1 = target_ls[i + 1]
            # print('~~~', i, target_label_0, target_label_1)

            target_lang, target_tag = target_i.split('_')[1:]

            pattern = context_dc.get(f'{target_lang}_info').get(target_tag)
            assert pattern, f'{(target_lang, target_tag)}不存在对应pattern?'

            # pattern = r'%r' % pattern
            match = re.search(pattern, replaced_template_text)
            if match:
                span = match.span()
                label_i[span[0]] = labels.index(target_label_0)
                for i in range(span[0] + 1, span[1]):
                    label_i[i] = labels.index(target_label_1)
                label_span.update({f'{target_lang}__{target_tag}': span})

        return label_i, label_span

    def get_replaced_template_text(self, template_file_path):
        context_dc = self.get_context_dc()
        replaced_template_text = template_replace(template_file_path, context_dc)

        label_i, label_span = self.get_label_i_and_span(replaced_template_text, context_dc)
        return replaced_template_text, label_i, label_span


if __name__ == '__main__':
    data_maker = DataMaker()
    # template_file_path_ls = [
    #     'templates/style_1.html',
    #     'templates/style_2.html',
    #     'templates/style_3.html',
    #     'templates/style_4.html',
    # ]

    template_file_name = 'style_3'
    template_file_path = f'templates/{template_file_name}.html'

    output_dir = f'tempdir/outputs/{template_file_name}'
    os.makedirs(output_dir, exist_ok=True)

    from bdtime import tt
    from tqdm import tqdm

    ls = []
    save_n = 1000

    total = 10000
    tq_i = tqdm(total=total)

    # for i in range(100):
    #     k = random.randint(0, 10)
    #     print(k)
    # exit()

    error_k = 0
    for i in range(total):
        try:
            replaced_template_text, label_i, label_span = data_maker.get_replaced_template_text(template_file_path)
            # print(label_span, '------', replaced_template_text)
            ls.append([replaced_template_text, label_i, label_span])
        except Exception as e:
            error_k += 1
            print(f'*** error {i}:', e)
            tq_i.desc = f"error_k: {error_k}, error_percent: {error_k / total * 100: .1f}%"
        if (i + 1) % save_n == 0:
            output_file_name = os.path.join(output_dir, f'{template_file_name}__n_{i + 1}.csv')

            print(f'---- save i: {i + 1}, len: {len(ls)}', ', output_file_name:', output_file_name)
            df = pd.DataFrame(ls, columns=data_maker.columns)
            df.to_csv(output_file_name, encoding='utf-8')
            ls = []
        tq_i.update(1)

    print(f'--- tt.now(): {tt.now()}')
